%% Evaluate the PITs

function [S_stat, S_p_value, PIT_plot] = PITs(hat_vartheta, bar_tau, predict_se, S_stat_sim, bin_num, ytick)
% INPUT: validation study estimates hat_vartheta, predictive distribution
% mean bar_tau, se predict_se, simulated sample of S statistics S_stat_sim,
% number of bins bin_num, ticks for y-axis in the plot
% OUTPUT: S statistics S_stat, p-value S_p_value

    N = size(hat_vartheta,1);

    % Initialization
    PIT = zeros(N, 1);
    
    % Compute the PIT for each element
    for i = 1:N
        PIT(i) = normcdf(hat_vartheta(i), bar_tau(i), predict_se(i));
    end
    
    % Define bins for the plot and S stat
    edges = linspace(0,1,bin_num+1);
    
    % Compute histogram counts for the bins
    [counts, ~] = histcounts(PIT, edges);
    
    % Calculate the S statistic
    expected_count = 0.2 * N;
    S_stat = sum((counts - expected_count).^2 / expected_count);
    
    % Calculate the p-value
    S_p_value = mean(S_stat_sim >= S_stat);

    % Plot the histogram of PITs
    frequencies = counts / N;
    PIT_plot = figure;
    set(PIT_plot, 'Units', 'inches', 'Position', [1 1 5.833 4.373]); 

    bar(linspace(1/(2*bin_num), 1 - 1/(2*bin_num), bin_num), frequencies);  % Plot at the midpoints of the bins
  
    xticks(edges); 
    xticklabels(arrayfun(@(x) sprintf('%.1f', x), edges, 'UniformOutput', false));  
    
    % Add a horizontal dashed line at y = 0.2
    hold on;
    yline(0.2, 'k--', 'LineWidth', 2.5);
    yticks(ytick);
    ylim([ytick(1), ytick(end)]);

    set(gca, 'FontSize', 18);
    hold off;

end